{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import pathlib\n",
    "import pickle\n",
    "\n",
    "from scapy.all import *  # noqa\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import random_split\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "from logger import logger\n",
    "from dataset import MTU_LENGTH\n",
    "\n",
    "dirname = pathlib.Path.cwd()\n",
    "plt.set_loglevel('info')\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    torch.set_default_tensor_type(torch.cuda.FloatTensor)\n",
    "    print(\"using cuda:\", torch.cuda.get_device_name(0))\n",
    "else:\n",
    "    torch.set_default_tensor_type(torch.FloatTensor)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PcapDataset(Dataset):\n",
    "\n",
    "    def __init__(self, vpn=False, type=False, app=False, test=False) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.vpn = vpn\n",
    "        self.type = type\n",
    "        self.app = app\n",
    "        self.test = test\n",
    "\n",
    "        with open(\"./data/packets.pickle\", 'rb') as file:\n",
    "            data = file.read()\n",
    "            self.packets = pickle.loads(data)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.packets)\n",
    "\n",
    "    def get_type_count(self):\n",
    "        if self.vpn:\n",
    "            return 2\n",
    "        if self.type:\n",
    "            return 6\n",
    "        if self.app:\n",
    "            return 16\n",
    "        raise Exception(\"unknown classify type\")\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        if index >= len(self.packets):\n",
    "            raise StopIteration\n",
    "\n",
    "        row = self.packets[index]\n",
    "\n",
    "        vpn, ty, app, content = row\n",
    "\n",
    "        target = torch.zeros(self.get_type_count())\n",
    "\n",
    "        if self.vpn:\n",
    "            label = vpn\n",
    "        elif self.type:\n",
    "            label = ty - 1\n",
    "        elif self.app:\n",
    "            label = app - 1\n",
    "\n",
    "        target[label] = 1.0\n",
    "\n",
    "        if len(content) < MTU_LENGTH:\n",
    "            content += b'\\0' * (MTU_LENGTH - len(content))\n",
    "\n",
    "        content = content[:MTU_LENGTH]\n",
    "\n",
    "        assert(len(content) == MTU_LENGTH)\n",
    "\n",
    "        data = np.frombuffer(content, dtype=np.uint8, count=MTU_LENGTH)\n",
    "\n",
    "        # image = torch.tensor(data.reshape(32, 32).copy()) / 255.0\n",
    "        image = torch.tensor(data.copy()) / 255.0\n",
    "        # image = image.view(1, 32, 32)\n",
    "        image = image.view(1, MTU_LENGTH)\n",
    "\n",
    "        return label, image, target\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Classifier(nn.Module):\n",
    "\n",
    "    def __init__(self, type_count):\n",
    "        super().__init__()\n",
    "\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv1d(1, 32, kernel_size=5, stride=1),\n",
    "            nn.LeakyReLU(0.02, inplace=True),\n",
    "            # # nn.ReLU(inplace=True),\n",
    "            # nn.BatchNorm2d(32),\n",
    "            nn.MaxPool1d(2, 2),\n",
    "\n",
    "            nn.Conv1d(32, 128, kernel_size=5, stride=2),\n",
    "            nn.LeakyReLU(0.02, inplace=True),\n",
    "            # # nn.ReLU(inplace=True),\n",
    "            # # nn.BatchNorm2d(64),\n",
    "            nn.MaxPool1d(2, 2),\n",
    "\n",
    "            nn.Conv1d(128, 64, kernel_size=5, stride=2),\n",
    "            nn.LeakyReLU(0.02, inplace=True),\n",
    "            # # nn.ReLU(inplace=True),\n",
    "            # # nn.BatchNorm2d(64),\n",
    "\n",
    "            nn.Flatten(),\n",
    "\n",
    "            # nn.LazyLinear(256),\n",
    "            # # # nn.Dropout(p=0.05),\n",
    "            # nn.LeakyReLU(0.02, inplace=True),\n",
    "\n",
    "            # # nn.LazyLinear(128),\n",
    "            # # nn.LeakyReLU(0.02, inplace=True),\n",
    "\n",
    "            nn.LazyLinear(64),\n",
    "            # # nn.Dropout(p=0.05),\n",
    "            nn.LeakyReLU(0.02, inplace=True),\n",
    "            # nn.ReLU(inplace=True),\n",
    "\n",
    "            nn.LazyLinear(type_count),\n",
    "            # # nn.Sigmoid(),\n",
    "            # # nn.LeakyReLU(0.02, inplace=True),\n",
    "        )\n",
    "\n",
    "        # self.loss_function = nn.MSELoss()\n",
    "        self.loss_function = nn.CrossEntropyLoss()\n",
    "\n",
    "        # self.optimizer = torch.optim.SGD(self.parameters(), lr=0.01)\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=0.01)\n",
    "\n",
    "        self.right = 0\n",
    "        self.total = 0\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        return self.model(inputs)\n",
    "\n",
    "    def train(self, label, inputs, targets):\n",
    "        outputs = self.forward(inputs)\n",
    "\n",
    "        loss = self.loss_function(outputs, targets)\n",
    "\n",
    "        self.optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        self.optimizer.step()\n",
    "\n",
    "        predicted = torch.max(outputs.data, 1)[1]\n",
    "\n",
    "        self.right += (predicted == label).sum()\n",
    "        self.total += len(label)\n",
    "\n",
    "        return loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kwargs = {\n",
    "    'vpn': False,\n",
    "    'type': False,\n",
    "    'app': True,\n",
    "}\n",
    "\n",
    "epoch = 64\n",
    "test_percent = 0.05\n",
    "\n",
    "writer = SummaryWriter()\n",
    "\n",
    "dataset = PcapDataset(**kwargs)\n",
    "c = Classifier(dataset.get_type_count())\n",
    "\n",
    "total_count = len(dataset)\n",
    "test_count = int(total_count * test_percent)\n",
    "train_count = total_count - test_count\n",
    "\n",
    "\n",
    "label, image, target = dataset[0]\n",
    "print(image.device, target.device)\n",
    "\n",
    "output = c.forward(image.view(1, 1, MTU_LENGTH))\n",
    "print(output.shape)\n",
    "\n",
    "trainset, testset = random_split(\n",
    "    dataset,\n",
    "    [train_count, test_count],\n",
    "    torch.Generator(device=device))\n",
    "\n",
    "loader = DataLoader(\n",
    "    dataset=trainset,\n",
    "    batch_size=512,\n",
    "    shuffle=True,\n",
    "    generator=torch.Generator(device=device),\n",
    "    drop_last=True,\n",
    ")\n",
    "\n",
    "for var in range(epoch):\n",
    "\n",
    "    tq = tqdm(loader)\n",
    "\n",
    "    c.total = 0\n",
    "    c.right = 0\n",
    "    for label, image, target in tq:\n",
    "        loss = c.train(label, image, target)\n",
    "        acc = c.right / c.total\n",
    "\n",
    "        writer.add_scalar(\"loss\", loss.item())\n",
    "        writer.add_scalar(\"acc\", acc)\n",
    "\n",
    "        tq.set_postfix(epoch=f\"{var}\", acc='%.6f' % acc)\n",
    "\n",
    "    tq.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 测试数据\n",
    "total = 0\n",
    "right = 0\n",
    "tq = tqdm(testset)\n",
    "\n",
    "rights = np.zeros(dataset.get_type_count(), dtype=np.float32)\n",
    "totals = np.zeros(dataset.get_type_count(), dtype=np.float32)\n",
    "\n",
    "for label, image, target in tq:\n",
    "    outputs = c.forward(image.view(1, 1, MTU_LENGTH))\n",
    "    total += 1\n",
    "    totals[label] += 1\n",
    "    if outputs.argmax() == label:\n",
    "        right += 1\n",
    "        rights[label] += 1\n",
    "\n",
    "    writer.add_scalar(\"tacc\", right / total)\n",
    "    tq.set_postfix(acc='%.6f' % (right / total))\n",
    "\n",
    "tq.close()\n",
    "\n",
    "print(right, total, right / total)\n",
    "\n",
    "print(totals)\n",
    "print(rights)\n",
    "print(rights / totals)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "9125777025588d260840c866d8b6bdfbdd8195897993998b0e02bb35dd13277b"
  },
  "kernelspec": {
   "display_name": "Python 3.7.2 ('test')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
